# -*- coding: utf-8 -*-
"""
@Time ： 2025/7/24 22:42
@Auth ： 韩明航
"""
import pandas as pd
from langchain.text_splitter import RecursiveCharacterTextSplitter

# 定义文件所在的路径
DOC_PATH = "成都.csv"

# 使用 pandas 加载 CSV 文件
df = pd.read_csv(DOC_PATH)

# 提取需要的列
columns_to_extract = ['景区特点描述', '景区榜单', '优待政策']
documents = df[columns_to_extract].values.tolist()

# 将提取的文本内容转换为字符串列表
text_data = ["\n".join(map(str, doc)) for doc in documents]

# 文本处理
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=100,  # 尝试调整它
    chunk_overlap=20,  # 尝试调整它
    separators=["\n\n", "\n", " ", "。", ""]
)

# 将文本分割成小块
docs = []
for text in text_data:
    chunks = text_splitter.split_text(text)
    docs.extend(chunks)

# 打印结果以验证
for doc in docs[:5]:
    # 打印前5个文本块
    print(len(doc))
    print(doc)
    print("------")

from langchain_huggingface import HuggingFaceEmbeddings

# 指定要加载的预训练模型的名称，参考排行榜：https://huggingface.co/spaces/mteb/leaderboard
model_name = "chuxin-llm/Chuxin-Embedding"

# 创建 Hugging Face 的嵌入模型实例，这个模型将用于将文本转换为向量表示（embedding）
embedding_model = HuggingFaceEmbeddings(model_name=model_name)

# 打印嵌入模型的配置信息，显示模型结构和其他相关参数
print(embedding_model)

# embed_query() 方法会将文本转换为嵌入的向量
query_embedding = embedding_model.embed_query("Hello")

# 打印生成的嵌入向量的长度，向量长度应与模型的输出维度一致（这里是 1024），你也可以选择打印向量看看
print(f"嵌入向量的维度为: {len(query_embedding)}")

from langchain.docstore.document import Document
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings

# 假设 embedding_model 已经定义
embedding_model = HuggingFaceEmbeddings(model_name="chuxin-llm/Chuxin-Embedding")

# 将字符串列表转换为 Document 对象列表
documents = [Document(page_content=doc) for doc in docs]

# 使用预训练嵌入模型生成向量并创建向量数据库
vectorstore = FAISS.from_documents(documents, embedding_model)

import numpy as np

# 获取向量数据
index = vectorstore.index
num_vectors = index.ntotal  # 获取向量总数
vector_dim = index.d  # 获取向量维度

# 打印向量总数和维度
print(f"向量总数: {num_vectors}")
print(f"向量维度: {vector_dim}")

# 获取向量数据
vectors = np.zeros((num_vectors, vector_dim), dtype=np.float32)
for i in range(num_vectors):
    vectors[i] = index.reconstruct(i)

print("向量数据:")
print(vectors)

# 获取元数据
docstore = vectorstore.docstore
documents = list(docstore._dict.values())  # 获取所有文档

# 打印元数据
print("元数据:")
for doc in documents:
    print(doc)
    print("------")

query = "适合和孩子一起玩的地方"

retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
# 检索与 query 相关的文档片段
retrieved_docs = retriever.invoke(query)

# 打印检索到的文档片段
for i, doc in enumerate(retrieved_docs):
    print(f"Document {i + 1}:")
    print(f"Content: {doc.page_content}\n")

# 使用 FAISS 数据库进行相似性搜索，返回最相关的文档片段
retrieved_docs = vectorstore.similarity_search_with_score(query, k=3)

# 现在的 retrieved_docs 包含 (Document, score)
for doc, score in retrieved_docs:
    print(f"Score: {score}")
    print(f"Content: {doc.page_content}\n")